#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Copyright (c) 61DUke.

# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
# See the Readme Document for more details.
# Create: 2020-12-27


from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D
from tensorflow.keras.models import Model
from sklearn.neighbors import NearestNeighbors
import numpy as np
from PIL import Image
import random

__all__ = ["retrieval_images_by_autoencoder"]


def retrieval_images_by_autoencoder(fit_images_path: list,
                                    target_images_path: list,
                                    scale=28,
                                    epochs=10,
                                    batch_size=32,
                                    n_neigh=5,
                                    shuffle=True
                                    ):
    """
    Build a very simple image retrieval system using a special type of Neural Network, called an autoencoder

    :param fit_images_path:  training images data path sequence
    :param target_images_path:  searching images data path sequence in training images data path sequence
    :param scale:  resize scale
    :param epochs:  training epochs, default 10
    :param batch_size: min-batch size, default 32
    :param n_neigh: nearest images number, default 5
    :param shuffle: training data random order
    :return: dict

    Example:
    ---------
        >>> from utils import get_file_paths

        >>> images_paths = list(get_file_paths(extension=".JPG", folder_path="data"))
        >>> nearest_images = retrieval_images_by_autoencoder(
        ...     fit_images_path=images_paths,
        ...     target_images_path=["122.JPG", "017.JPG"],
        ...     scale=28,
        ...     epochs=15,
        ...     batch_size=32,
        ...     n_neigh=5,
        ...     shuffle=True
        >>> )
        >>> print(nearest_images)
        >>> {
        ...    "122.JPG": [
        ...        (1.0, "data/124.JPG"),
        ...        (0.021841750088447576, "data/535.JPG"),
        ...        (0.013354248283854897, "data/605.JPG"),
        ...        (0.009994627751588991, "data/402.JPG"),
        ...        (0.0, "data/134.JPG")
        ...    ],
        ...    "017.JPG": [
        ...        (0.11133096949872578, "data/510.JPG"),
        ...        (0.03817764172216209, "data/471.JPG"),
        ...        (0.02644454978659949, "data/IMG_2044.JPG"),
        ...        (0.001042320143904596, "data/130.JPG"),
        ...        (0.0, "data/472.JPG")
        ...    ]
        >>> }
    """
    if not isinstance(fit_images_path, list) or not fit_images_path:
        raise ValueError("fit_images_path param must be list!")
    if not isinstance(target_images_path, list) or not target_images_path:
        raise ValueError("target_images_path param must be list!")
    if not isinstance(scale, int) or not scale:
        raise ValueError("target_images_path param must be int!")
    if not isinstance(epochs, int) or not epochs:
        raise ValueError("target_images_path param must be int!")
    if not isinstance(batch_size, int) or not batch_size:
        raise ValueError("target_images_path param must be int!")
    if not isinstance(n_neigh, int) or not n_neigh:
        raise ValueError("target_images_path param must be int!")
    if not isinstance(shuffle, bool):
        raise ValueError("shuffle param must be bool!")

    if shuffle:
        random.shuffle(fit_images_path)

    # Gray transformation of image, Reshape the data to have 1 channel
    X = np.array(list(map(lambda img: np.array(
        Image.open(img).convert("L").resize((scale, scale), Image.ANTIALIAS)
    ), fit_images_path))).astype("float32") / 255.0

    target_images = np.array(list(map(lambda img: np.array(
        Image.open(img).convert("L").resize((scale, scale), Image.ANTIALIAS)
    ), target_images_path))).astype("float32") / 255.0

    X = np.reshape(X, (-1, scale, scale, 1))
    target_images = np.reshape(target_images, (-1, scale, scale, 1))

    # Convolution, Create the autoencoder
    input_images = Input(shape=(scale, scale, 1))
    x = Conv2D(16, (3, 3), activation="relu", padding="same")(input_images)
    x = MaxPooling2D((2, 2), padding="same")(x)
    x = Conv2D(8, (3, 3), activation="relu", padding="same")(x)
    x = MaxPooling2D((2, 2), padding="same")(x)
    x = Conv2D(8, (3, 3), activation="relu", padding="same")(x)
    encoded = MaxPooling2D((2, 2), padding="same", name="encoder")(x)

    x = Conv2D(8, (3, 3), activation="relu", padding="same")(encoded)
    x = UpSampling2D((2, 2))(x)
    x = Conv2D(8, (3, 3), activation="relu", padding="same")(x)
    x = UpSampling2D((2, 2))(x)
    x = Conv2D(16, (3, 3), activation="relu")(x)
    x = UpSampling2D((2, 2))(x)
    decoded = Conv2D(1, (3, 3), activation="sigmoid", padding="same")(x)

    autoencoder = Model(input_images, decoded)
    autoencoder.compile(optimizer="adam", loss="mse")
    # Train it
    autoencoder.fit(X, X, epochs=epochs, batch_size=batch_size, callbacks=None)
    autoencoder.summary()

    # Create the encoder part
    # The encoder part is the first half of the autoencoder, i.e.
    # the part that will encode the input into a latent space representation.
    encoder = Model(inputs=autoencoder.input, outputs=autoencoder.get_layer("encoder").output)

    codes = encoder.predict(X)
    query_codes = encoder.predict(target_images)

    # Encode the query image and the test set
    codes = codes.reshape(-1, int(np.ceil(scale / 2 ** 3)) * int(np.ceil(scale / 2 ** 3)) * 8)
    query_codes = query_codes.reshape(-1, int(np.ceil(scale / 2 ** 3)) * int(np.ceil(scale / 2 ** 3)) * 8)
    # Find the closest images
    nbrs = NearestNeighbors(n_neighbors=n_neigh).fit(codes)
    distances, indices = nbrs.kneighbors(query_codes)

    target_images_nearest = {
        target_images_path[i]: list(zip(
            map(lambda d: 1.0 - (d / distance[-1]), distance),
            map(lambda ind_: fit_images_path[ind_], ind)
        )) for i, (distance, ind) in enumerate(zip(distances, indices))}

    return target_images_nearest


if __name__ == '__main__':
    pass
